﻿#ifndef _CELL_IOCP_HPP_
#define _CELL_IOCP_HPP_

#ifdef _WIN32

#ifndef CELL_USE_IOCP
#define CELL_USE_IOCP
#endif // #ifndef CELL_USE_IOCP

#include "CELL.h"
#include <MSWSock.h>
#include <stdio.h>

enum IO_TYPE
{
	ACCEPT = 10,
	RECV,
	SEND
};

// 数据缓冲区大小
#define DATA_BUFF_SIZE 1024

struct IO_DATA_BASE
{
	// 重叠体
	OVERLAPPED overlapped;
	// 
	SOCKET sockfd;
	/*
		数据缓冲区
		? 为什么要这样一个结构存储数据，数据指针和记录指向数据的大小
		!: 每个客户端都有自己的接收和发送缓冲区，那么在定义多余的空间就有点浪费，
			所以只要一个指针指向接收或者发送缓冲区数据具体的位置就可以，
			IOCP在接收数据完成以后就会通知接收或发送数据了，而不是跟
			select和epoll一样通知可以接收/发送才告诉程序，IOCP是接收完/发送完
			才会告知程序
	*/
	WSABUF wsaBuff;
	// 操作类型
	IO_TYPE iotype;
};

struct IO_EVENT
{
	union
	{
		void* ptr;
		SOCKET sockfd;
	} data;
	IO_DATA_BASE* pIOData;
	DWORD bytesTrans = 0;
};

class CELLIOCP
{
public:
	CELLIOCP()
	{

	}

	~CELLIOCP()
	{
		destory();
	}

	// 将AcceptEX函数加载到内存中
	bool loadAcceptEX(SOCKET ListenSocket)
	{
		if (INVALID_SOCKET != _sockServer)
		{
			CELLLog_Error("loadAcceptEX _sockServer != INVALID_SOCKET");
			return false;
		}
		if (_AcceptEx)
		{
			CELLLog_Error("loadAcceptEX _AcceptEx != NULL");
			return false;
		}
		_sockServer = ListenSocket;
		GUID GuidAcceptEx = WSAID_ACCEPTEX;
		DWORD dwBytes = 0;
		int iResult = WSAIoctl(ListenSocket, SIO_GET_EXTENSION_FUNCTION_POINTER,
			&GuidAcceptEx, sizeof(GuidAcceptEx),
			&_AcceptEx, sizeof(_AcceptEx),
			&dwBytes, NULL, NULL);
		if (iResult == SOCKET_ERROR) {
			CELLLog_Error("WSAIoctl failed with error: %u", WSAGetLastError());
			return false;
		}
		return true;
	}

	// 创建一个IO完成端口(IOCP)
	int create()
	{
		_completionPort = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0);
		if (NULL == _completionPort)
		{
			CELLLog_Error("CELLIOCP create failed with error %d", GetLastError());
			return -1;
		}
		return 0;
	}

	// 销毁IOCP
	void destory()
	{
		if (_completionPort)
		{
			CloseHandle(_completionPort);
			_completionPort = NULL;
		}
	}

	// 关联文件sockfd和IOCP
	HANDLE reg(SOCKET sockfd)
	{
		auto ret = CreateIoCompletionPort((HANDLE)sockfd, _completionPort, (ULONG_PTR)sockfd, 0);
		if (NULL == ret)
		{
			CELLLog_Error("CELLIOCP reg sockfd failed with error %d", GetLastError());
			return NULL;
		}
		return ret;
	}

	// 关联自定义数据地址和IOCP
	HANDLE reg(SOCKET sockfd, void* ptr)
	{
		auto ret = CreateIoCompletionPort((HANDLE)sockfd, _completionPort, (ULONG_PTR)ptr, 0);
		if (NULL == ret)
		{
			CELLLog_Error("CELLIOCP reg ptr failed with error %d", GetLastError());
			return NULL;
		}
		return ret;
	}

	// 投递接收链接任务
	bool postAccept(IO_DATA_BASE* pIoData)
	{
		if (!_AcceptEx)
		{
			CELLLog_Error("postAccept _AcceptEX is NULL");
			return false;
		}
		pIoData->iotype = IO_TYPE::ACCEPT;
		pIoData->sockfd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
		/*
			! AcceptEx函数第四个参数只要是非0值，那么只有连接的客户端发送
				一条消息(大小随意)后，IOCP才会告诉程序这个客户端连接了
			? 为什么第三个参数这样写
				ack: 远端地址和本地地址都会存储在buffer尾部，
					而远端地址和本地地址的大小都是sizeof(sockaddr_in) + 16,
					剩下960字节可用
		*/
		if (FALSE == _AcceptEx(_sockServer,
			pIoData->sockfd,
			pIoData->wsaBuff.buf,
			// sizeof(ioData.buffer) - (sizeof(sockaddr_in) + 16) * 2,
			0,
			sizeof(sockaddr_in) + 16,
			sizeof(sockaddr_in) + 16,
			NULL,
			&pIoData->overlapped))
		{
			int errCode = WSAGetLastError();
			// ERROR_IO_PENDING:操作还在进行，没有出错
			if (errCode != ERROR_IO_PENDING)
			{
				/*
					在打开了超过接收上限的连接再关闭后，再关掉一个现有连接，
						此时再投递接收数据，会失败并报错以下错误
						WSAECONNRESET                    10054L
				*/
				if (WSAECONNRESET == errCode)
				{
					// 远端关闭
					return false;
				}
				CELLLog_Error("acceptexRet falied with error %d", errCode);
				return false;
			}
		}
		return true;
	}

	// 投递接收数据任务
	bool postRecv(IO_DATA_BASE* pIoData)
	{
		pIoData->iotype = IO_TYPE::RECV;
		DWORD flags = 0;
		ZeroMemory(&pIoData->overlapped, sizeof(OVERLAPPED));

		if (SOCKET_ERROR == WSARecv(pIoData->sockfd, &pIoData->wsaBuff, 1, NULL, &flags, &pIoData->overlapped, NULL))
		{
			int err = WSAGetLastError();
			if (ERROR_IO_PENDING != err)
			{
				if (WSAECONNRESET == err)
				{
					// 远端客户端关闭
					return false;
				}
				CELLLog_Error("CELLIOCP WSARecv failed with error %d", err);
				return false;
			}
		}
		return true;
	}

	// 投递发送数据任务
	bool postSend(IO_DATA_BASE* pIoData)
	{
		pIoData->iotype = IO_TYPE::SEND;
		DWORD flags = 0;
		ZeroMemory(&pIoData->overlapped, sizeof(OVERLAPPED));

		if (SOCKET_ERROR == WSASend(pIoData->sockfd, &pIoData->wsaBuff, 1, 0, flags, &pIoData->overlapped, NULL))
		{
			int err = WSAGetLastError();
			if (ERROR_IO_PENDING != err)
			{
				if (WSAECONNRESET == err)
				{
					// 远端关闭
					return false;
				}
				printf("CELLIOCP WSASend failed with error %d\n", err);
				return false;
			}
		}
		return true;
	}

	// 获取所有任务的状态
	int wait(IO_EVENT& ioEvent, unsigned int timeout = INFINITE)
	{
		ioEvent.bytesTrans = 0;
		ioEvent.pIOData = NULL;
		ioEvent.data.ptr = NULL;
		// 最后参数如果传INFINITE会等到有事件完成才返回(阻塞)
		if (FALSE == GetQueuedCompletionStatus(_completionPort,
			&ioEvent.bytesTrans,
			(PULONG_PTR)& ioEvent.data,
			(LPOVERLAPPED*)&ioEvent.pIOData,
			timeout))
		{
			int err = GetLastError();
			if (WAIT_TIMEOUT == err)
			{
				// 超时
				return 0;
			}
			else if (ERROR_NETNAME_DELETED == err)
			{
				// 客户端断开连接
				return 1;
			}
			else if (ERROR_CONNECTION_ABORTED == err)
			{
				// 主动断开客户端连接
				return 1;
			}
			CELLLog_Error(" CELLIOCP GetQueuedCompletionStatus falied with error %d", err);
			return -1;
		}
		return 1;
	}
private:
	// 接收客户端函数指针
	LPFN_ACCEPTEX _AcceptEx = NULL;
	// IOCP完成端口
	HANDLE _completionPort = NULL;
	// 
	SOCKET _sockServer = INVALID_SOCKET;
};

#endif // #ifdef _WIN32

#endif // #ifndef _CELL_IOCP_HPP_
